package com.elitel;

import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.util.UriComponentsBuilder;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.nio.charset.StandardCharsets;

@Component
@Slf4j
public class ReqRespDealFilter implements GlobalFilter, Ordered {

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        String rawPath = request.getURI().getRawPath();
        MultiValueMap<String, String> queryParams = request.getQueryParams();
        Flux<DataBuffer> body = request.getBody();
        return body
                .collectList()
                .flatMap(dataBuffers -> {
                    // 合并所有DataBuffer为一个字符串（请求体）
                    StringBuilder bodyBuilder = new StringBuilder();
                    dataBuffers.forEach(buffer -> {
                        byte[] bytes = new byte[buffer.readableByteCount()];
                        buffer.read(bytes);
                        bodyBuilder.append(new String(bytes));
                        DataBufferUtils.release(buffer); // 释放缓冲区
                    });
                    // 原始请求体
                    String originalBodyStr = bodyBuilder.toString();

                    // 2. 处理请求体（如果需要修改）
                    String modifiedBodyStr = originalBodyStr; // 这里可以添加修改逻辑

                    // 3. 创建新的请求体
                    DataBufferFactory bufferFactory = exchange.getResponse().bufferFactory();
                    DataBuffer modifiedBody = bufferFactory.wrap(modifiedBodyStr.getBytes());

                    MultiValueMap<String, String> newQueryParams =new LinkedMultiValueMap<>();
                    newQueryParams.add("cc", "dd");
                    newQueryParams.addAll(queryParams);
                    URI uri = UriComponentsBuilder.fromUri(request.getURI())
                            .replaceQueryParams(newQueryParams).build().toUri();

                    // 4. 装饰请求对象
                    ServerHttpRequest modifiedRequest = new ServerHttpRequestDecorator(request) {
                        @Override
                        public Flux<DataBuffer> getBody() {
                            return Flux.just(modifiedBody);
                        }

                        @Override
                        public HttpHeaders getHeaders() {
                            HttpHeaders headers = new HttpHeaders();
                            headers.putAll(request.getHeaders());
                            return headers;
                        }


                        @Override
                        public MultiValueMap<String, String> getQueryParams() {
                            return queryParams;
                        }

                        @Override
                        public URI getURI() {
                            return uri;
                        }
                    };
                    // 6. 处理响应
                    ServerHttpResponse originalResponse = exchange.getResponse();
                    // 装饰响应对象以获取响应体
                    ServerHttpResponseDecorator decoratedResponse = new ServerHttpResponseDecorator(originalResponse) {
                        public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                            if (body instanceof Flux) {
                                Flux<? extends DataBuffer> fluxBody = Flux.from(body);
                                return super.writeWith(fluxBody.collectList().flatMap(dataBuffers -> {
                                    StringBuilder bodyBuilder = new StringBuilder();
                                    dataBuffers.forEach(buffer -> {
                                        byte[] bytes = new byte[buffer.readableByteCount()];
                                        buffer.read(bytes);
                                        bodyBuilder.append(new String(bytes, StandardCharsets.UTF_8));
                                        // 释放缓冲区
                                        DataBufferUtils.release(buffer);
                                    });
                                    String string = bodyBuilder.toString();

                                    return Mono.just(bufferFactory.wrap(string.getBytes(StandardCharsets.UTF_8)));
                                }));

                            }
                            return super.writeWith(body);
                        }
                    };

                    // 7. 继续过滤器链
                    return chain.filter(exchange.mutate()
                            .request(modifiedRequest)
                            .response(decoratedResponse)
                            .build());
                });
    }

    @Override
    public int getOrder() {
        return -3;
    }

}